import os
import json,re
from ollama import chat
from ollama import ChatResponse
def clean_ollama_response(response_text):

    # 使用正则表达式清除代码块标记和换行符
    cleaned_text = re.sub(r'```json\n|```|\n', '', response_text.replace("\\",""))
    return cleaned_text

def send_to_ollama(file_path, model='qwen2.5:14b'):
    # 读取文件内容
    with open(file_path, 'r', encoding='utf-8') as file:
        content = file.read()
    
    # 构建请求的messages
    response: ChatResponse = chat(model=model, messages=[{
        'role': 'user',
        'content': content+'-->请提取上面文本中所有吕小鱼的问答，将她的心理活动也尽量转换为问答，每个问答返回json格式的数据:"{"input": 问句,"output": 答句},"输出为标准正确的json格式数据,不要回答其他任何多余的内容和多余的字符或结构，如果问答不完整请根据吕小鱼风格补全，如何没有有关吕小鱼的问答返回null',
    }])
    
    # 获取模型输出
    response_text = response.message.content
    return response_text

def save_responses_to_file(all_responses, output_dir, part_number):
    # 保存每100个文件处理后的结果到一个新的文件
    output_file = os.path.join(output_dir, f'training_data_part_{part_number}.json')
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(all_responses, f, ensure_ascii=False, indent=4)
    print(f"Processed {len(all_responses)} responses and saved to {output_file}")

def process_files(input_dir='sp', output_dir='output', model='qwen2.5:14b', batch_size=30):
    # 确保输出文件夹存在
    os.makedirs(output_dir, exist_ok=True)

    all_responses = []
    processed_count = 0
    part_number = 1

    # 遍历目录中的文件
    for filename in sorted(os.listdir(input_dir)):
        if filename.endswith('.txt'):
            file_path = os.path.join(input_dir, filename)
            print(f"Processing file: {file_path}")
            response_text = send_to_ollama(file_path, model)
            
            if response_text:
                clean_response_text=clean_ollama_response(response_text)
                if clean_response_text:
                    all_responses.append(clean_response_text)
                processed_count += 1

                    # 每处理完100个文件保存一次
                if processed_count >= batch_size:
                    save_responses_to_file(all_responses, output_dir, part_number)
                    all_responses = []  # 清空当前批次的响应数据
                    processed_count = 0  # 重置计数器
                    part_number += 1  # 增加文件编号

    # 如果还有未保存的响应，保存到最后的文件
    if all_responses:
        save_responses_to_file(all_responses, output_dir, part_number)

# 使用示例
process_files(input_dir='sp', output_dir='output')
